﻿#pragma once

#include <algorithm>
#include <iostream>
#include <vector>

#include <xy/math/math.h>

namespace xy
{

namespace math
{

/**
 * @brief
 *
 * @tparam T 系数类型
 * @tparam Real 参数类型
 */
template <class T, class Real = T> class poly
{
  protected:
    static constexpr auto cmp = compare{};

    Real m_offset{};
    std::vector<T> m_coefficients;

  public:
    poly(const poly &) = default;
    poly(poly &&) = default;
    poly &operator=(const poly &) = default;
    poly &operator=(poly &&) = default;

    /**
     * @brief 限定显式调用，防止 operator* 中的隐式转换
     *
     * @param[in] deg
     */
    explicit poly(std::size_t deg) : m_coefficients(deg + 1, T{})
    {
    }

    /**
     * @brief 系数构造函数
     *
     * @param[in] coefficients
     */
    explicit poly(std::vector<T> &&coefficients) : m_coefficients(std::move(coefficients))
    {
    }

    /**
     * @brief 比较
     *
     * @param[in] poly
     * @return true
     * @return false
     */
    bool operator==(const poly &inPoly) const
    {
        if (degree() != inPoly.degree())
            return false;

        for (std::size_t i = 0; i < degree() + 1; i++)
            if (!cmp.less<Real>(m_coefficients[i], inPoly.m_coefficients[i]))
                return false;

        return true;
    }

    /**
     * @brief 参数偏移
     *
     * @param[in] offset
     * @return poly
     */
    poly &operator>>(Real offset)
    {
        m_offset += offset;
        return *this;
    }

    /**
     * @brief 获得偏移值
     *
     * @return Real
     */
    Real offset() const
    {
        return m_offset;
    }

    /**
     * @brief 获得次数
     *
     * @return std::size_t
     */
    std::size_t degree() const noexcept
    {
        return m_coefficients.size() - 1;
    }

    /**
     * @brief 升次
     *
     * @param[in] d
     * @return true
     * @return false
     */
    bool elevate(std::size_t d)
    {
        if (d <= degree())
            return false;

        // 高位填 0
        m_coefficients.resize(d + 1, T{});
        return true;
    }

    /**
     * @brief 多项式降阶
     *
     * @param[in] d
     * @return true
     * @return false
     */
    bool reduce(std::size_t d)
    {
        if (d >= degree())
            return false;

        // 低位截断
        m_coefficients.resize(d + 1, T{});
        return true;
    }

    /**
     * @brief 获得系数
     *
     * @return const std::vector<T>&
     */
    const std::vector<T> &coefficients() const noexcept
    {
        return m_coefficients;
    }

    /**
     * @brief 引用系数
     *
     * @return std::vector<T>&
     */
    std::vector<T> &coefficients() noexcept
    {
        return m_coefficients;
    }

    /**
     * @brief 获得系数
     *
     * @param[in] i
     * @return T
     */
    T coefficient(std::size_t i) const
    {
        return m_coefficients[i];
    }

    /**
     * @brief 引用系数
     *
     * @param[in] i
     * @return T&
     */
    T &coefficient(std::size_t i)
    {
        return m_coefficients[i];
    }

#define POLY_OP(func, op)                                                                                              \
    poly func(const poly &inPoly) const                                                                                \
    {                                                                                                                  \
        poly ret(std::max(degree(), inPoly.degree()));                                                                 \
        auto &coeff = ret.coefficients();                                                                              \
        const auto &coeff1 = m_coefficients;                                                                           \
        const auto &coeff2 = inPoly.coefficients();                                                                    \
        for (std::size_t i = 0; i < coeff1.size(); i++)                                                                \
            coeff[i] = coeff1[i];                                                                                      \
        for (std::size_t i = 0; i < coeff2.size(); i++)                                                                \
            coeff[i] = coeff[i] op coeff2[i];                                                                          \
        return ret;                                                                                                    \
    }

    POLY_OP(operator+, +)
    POLY_OP(operator-, -)
#undef POLY_OP

#define POLY_OP(func, op)                                                                                              \
    const poly &func(const poly &inPoly)                                                                               \
    {                                                                                                                  \
        elevate(inPoly.degree());                                                                                      \
        auto &coeff = m_coefficients;                                                                                  \
        const auto &coeff2 = inPoly.coefficients();                                                                    \
        for (std::size_t i = 0; i < coeff2.size(); i++)                                                                \
            coeff[i] op coeff2[i];                                                                                     \
        return *this;                                                                                                  \
    }

    POLY_OP(operator+=, +=)
    POLY_OP(operator-=, -=)
#undef POLY_OP

#define POLY_OP(func, op)                                                                                              \
    poly func(T k) const                                                                                               \
    {                                                                                                                  \
        poly ret(degree());                                                                                            \
        for (std::size_t i = 0; i < m_coefficients.size(); i++)                                                        \
            ret.m_coefficients[i] = m_coefficients[i] op k;                                                            \
        return ret;                                                                                                    \
    }

    POLY_OP(operator*, *)
    POLY_OP(operator/, /)
#undef POLY_OP

#define POLY_OP(func, op)                                                                                              \
    poly &func(T k)                                                                                                    \
    {                                                                                                                  \
        for (std::size_t i = 0; i < m_coefficients.size(); i++)                                                        \
            m_coefficients[i] op k;                                                                                    \
        return *this;                                                                                                  \
    }

    POLY_OP(operator*=, *=)
    POLY_OP(operator/=, /=)
#undef POLY_OP

    /**
     * @brief 多项式乘法
     *
     * @param[in] inPoly
     * @return poly
     */
    poly operator*(const poly &inPoly) const
    {
        poly ret(degree() + inPoly.degree());
        for (std::size_t i = 0; i < degree() + 1; i++)
            for (std::size_t j = 0; j < inPoly.degree() + 1; j++)
                ret.m_coefficients[i + j] += m_coefficients[i] * inPoly.m_coefficients[j];
        return ret;
    }

    /**
     * @brief 多项式乘法
     *
     * @param[in] inPoly
     * @return poly&
     */
    poly &operator*=(const poly &inPoly)
    {
        *this = *this * inPoly;
        return *this;
    }

    /**
     * @brief 多项式除法
     *
     * @param[in] inPoly
     * @return poly
     */
    poly operator/(const poly &inPoly) const
    {
        // 次数不够，返回 0 多项式（这里是唯一允许次数为负的情况）
        if (degree() < inPoly.degree())
            return poly(-1);

        poly ret(degree() - inPoly.degree());
        ret.m_coefficients[degree() - inPoly.degree()] = m_coefficients.back() / inPoly.m_coefficients.back();

        // 获得余式
        poly tmp = *this - ret * inPoly;

        // 减去最高项后需要降阶，然后递归计算商
        tmp.reduce(degree() - 1);
        tmp = tmp / inPoly;

        // 商式赋值
        for (std::size_t i = 0; i < tmp.degree() + 1; ++i)
            ret.m_coefficients[i] = tmp.m_coefficients[i];

        return ret;
    }

    /**
     * @brief 多项式除法
     *
     * @param[in] inPoly
     * @return poly&
     */
    const poly &operator/=(const poly &inPoly)
    {
        *this = *this / inPoly;
        return *this;
    }

    /**
     * @brief 求值
     *
     * @param[in] x
     * @return T
     */
    T value(Real x) const
    {
        std::size_t O = degree();
        if (O < 0)
            return T{};

        // 获得最高次系数，使用累乘方法提高效率
        T y = m_coefficients[O];
        x -= m_offset;
        for (int i = O - 1; i >= 0; i--)
            y = y * x + m_coefficients[i];

        return y;
    }

    /**
     * @brief 多项式定积分
     *
     * @param[in] a
     * @param[in] b
     * @return T
     */
    T integrate(Real a, Real b) const
    {
        T I = T{};
        a -= m_offset;
        b -= m_offset;
        for (std::size_t i = 0; i < degree() + 1; i++)
            I += m_coefficients[i] * (std::pow(b, i + 1) - std::pow(a, i + 1)) / (i + 1);
        return I;
    }

    /**
     * @brief 多项式不定积分
     *
     * @return poly
     */
    poly integrate() const
    {
        poly inPoly(degree() + 1);
        for (std::size_t i = 0; i < degree() + 1; i++)
            inPoly.coefficients()[i + 1] = m_coefficients[i] / (i + 1);
        return inPoly >> m_offset;
    }

    /**
     * @brief 多项式求导
     *
     * @return poly
     */
    poly derivative() const
    {
        if (degree() == 0)
            return poly(0);

        poly inPoly(degree() - 1);
        for (std::size_t i = 1; i < degree() + 1; i++)
            inPoly.coefficients()[i - 1] = m_coefficients[i] * i;
        return inPoly >> m_offset;
    }

    /**
     * @brief 产生节点多项式
     *
     * @param[in] c
     * @return poly
     */
    static poly nodes(const std::vector<Real> &c)
    {
        // 直接写 {1} 似乎会被错误识别
        poly inPoly(0);
        inPoly.coefficient(0) = T{1};
        for (std::size_t i = 0; i < c.size(); i++)
            inPoly = inPoly * poly({T{-c[i]}, T{1}});
        return inPoly;
    }

    /**
     * @brief 已知节点 c0 c1 ... cn ，获得 cj 对应的初等多项式
     *
     * @param[in] c
     * @param[in] j
     * @return poly
     */
    static poly lagrange(const std::vector<Real> &c, std::size_t j)
    {
        // 直接写 {1} 似乎会被错误识别
        poly inPoly(0);
        inPoly.coefficient(0) = T{1};
        for (std::size_t k = 0; k < c.size(); k++)
            if (k != j)
                inPoly = inPoly * poly({T{-c[k]}, T{1}}) / (c[j] - c[k]);
        return inPoly;
    }

    /**
     * @brief (-1,1) 上的 Gauss-Legendre 多项式，用于数值积分
     * 三项递推式 i * p_{i-1} - (2 * i + 1) * x * p_i + (i + 1) * p_{i+1} = 0
     *
     * @param[in] s
     * @return poly
     */
    static poly legendre(std::size_t s)
    {
        // 直接写 {1} 似乎会被错误识别
        poly f0(0);
        f0.coefficient(0) = T{1};
        if (s > 0)
        {
            poly x({T{0}, T{1}});
            poly f1 = x, f2 = x;
            for (std::size_t i = 1; i < s; i++)
            {
                f2 = x * f1 * (2 * i + 1) / (i + 1) - f0 * i / (i + 1);
                f0 = f1;
                f1 = f2;
            }
            f0 = f2;
        }
        return f0;
    }

    /**
     * @brief (0,1) 上的首一平移 Gauss-Legendre 多项式，用于 Runge-Kutta 积分
     *
     * @param[in] s
     * @return poly
     */
    static poly monic_legendre(std::size_t s)
    {
        // 注意不要使用 0 多项式
        poly inPoly(0);
        for (std::size_t j = 0; j <= s; j++)
        {
            poly v(j);
            v.coefficient(j) = T{Real(((s - j) % 2 == 0 ? 1 : -1))} * binomial(s, j) * binomial(s + j, j);
            inPoly = inPoly + v;
        }
        inPoly *= std::pow(factorial(s), 2) / factorial(2 * s);
        return inPoly;
    }
};

template <class T, class Real = T> class ppoly
{
  protected:
    static constexpr auto cmp = compare{};

    std::vector<Real> m_interval;       // 分段区间端点，区间左闭右开
    std::vector<poly<T, Real>> m_polys; // 多项式列表，数量是端点数 +1，包括左右半开区间上的多项式

  public:
    ppoly(const ppoly &) = default;
    ppoly(ppoly &&) = default;
    ppoly &operator=(const ppoly &) = default;
    ppoly &operator=(ppoly &&) = default;

    /**
     * @brief 空构造
     *
     */
    explicit ppoly() : m_polys({poly<T, Real>(0)})
    {
    }

    /**
     * @brief 单多项式
     *
     * @param[in] inPoly
     */
    explicit ppoly(poly<T, Real> &&inPoly) : m_polys({std::move(inPoly)})
    {
    }

    /**
     * @brief 只设置区间，每段区间上是零
     *
     * @param[in] interval
     */
    explicit ppoly(std::vector<Real> &&interval)
        : m_interval(std::move(interval)), m_polys(m_interval.size() + 1, poly<T, Real>(0))
    {
    }

    /**
     * @brief 设置区间和多项式
     *
     * @param[in] interval
     * @param[in] inPoly
     */
    ppoly(std::vector<Real> &&interval, std::vector<poly<T, Real>> &&inPoly)
        : m_interval(std::move(interval)), m_polys(std::move(inPoly))
    {
        // 检查赋值是否有效
        if (m_interval.size() + 1 != m_polys.size())
            throw std::logic_error("ppoly: invalid input size");
    }

    /**
     * @brief 比较
     *
     * @param[in] poly
     * @return true
     * @return false
     */
    bool operator==(const ppoly &inPoly) const
    {
        return m_interval == inPoly.m_interval && m_polys == inPoly.m_polys;
    }

#define POLY_OP(func, op)                                                                                              \
    ppoly func(const ppoly &inPoly) const                                                                              \
    {                                                                                                                  \
        if (m_interval.size() != inPoly.m_interval.size())                                                             \
            throw std::logic_error("ppoly: invalid input size");                                                       \
        for (std::size_t i = 0; i < m_interval.size(); ++i)                                                            \
            if (std::abs(m_interval[i] - inPoly.m_interval[i]) > std::numeric_limits<Real>::epsilon())                 \
                throw std::logic_error("ppoly: invalid input interval");                                               \
        auto interval = m_interval;                                                                                    \
        ppoly ret(std::move(interval));                                                                                \
        for (std::size_t i = 0; i < m_polys.size(); ++i)                                                               \
            ret.m_polys[i] = m_polys[i] op inPoly.m_polys[i];                                                          \
        return ret;                                                                                                    \
    }

    POLY_OP(operator+, +)
    POLY_OP(operator-, -)
    POLY_OP(operator*, *)
    POLY_OP(operator/, /)
#undef POLY_OP

#define POLY_OP(func, op)                                                                                              \
    ppoly &func(const ppoly &inPoly)                                                                                   \
    {                                                                                                                  \
        if (m_interval.size() != inPoly.m_interval.size())                                                             \
            throw std::logic_error("ppoly: invalid input size");                                                       \
        for (std::size_t i = 0; i < m_interval.size(); ++i)                                                            \
            if (std::abs(m_interval[i] - inPoly.m_interval[i]) > std::numeric_limits<Real>::epsilon())                 \
                throw std::logic_error("ppoly: invalid input interval");                                               \
        for (std::size_t i = 0; i < m_polys.size(); ++i)                                                               \
            m_polys[i] op inPoly.m_polys[i];                                                                           \
        return *this;                                                                                                  \
    }

    POLY_OP(operator+=, +=)
    POLY_OP(operator-=, -=)
    POLY_OP(operator*=, *=)
    POLY_OP(operator/=, /=)
#undef POLY_OP

#define POLY_OP(func, op)                                                                                              \
    ppoly func(T k) const                                                                                              \
    {                                                                                                                  \
        auto interval = m_interval;                                                                                    \
        ppoly ret(std::move(interval));                                                                                \
        for (std::size_t i = 0; i < m_polys.size(); ++i)                                                               \
            ret.m_polys[i] = m_polys[i] op k;                                                                          \
        return ret;                                                                                                    \
    }

    POLY_OP(operator*, *)
    POLY_OP(operator/, /)
#undef POLY_OP

#define POLY_OP(func, op)                                                                                              \
    ppoly &func(T k)                                                                                                   \
    {                                                                                                                  \
        for (auto &inPoly : m_polys)                                                                                   \
            inPoly op k;                                                                                               \
        return *this;                                                                                                  \
    }

    POLY_OP(operator*=, *=)
    POLY_OP(operator/=, /=)
#undef POLY_OP

    /**
     * @brief 获得多项式数量
     *
     * @return std::size_t
     */
    std::size_t poly_num() const noexcept
    {
        return m_polys.size();
    }

    /**
     * @brief 获得端点
     *
     * @return const std::vector<Real>&
     */
    const std::vector<Real> &intervals() const noexcept
    {
        return m_interval;
    }

    /**
     * @brief 获得多项式
     *
     * @return const std::vector<poly<T, Real>>&
     */
    const std::vector<poly<T, Real>> &polys() const noexcept
    {
        return m_polys;
    }

    /**
     * @brief 获得多项式
     *
     * @return std::vector<poly<T, Real>>&
     */
    std::vector<poly<T, Real>> &polys() noexcept
    {
        return m_polys;
    }

    /**
     * @brief 获得指定索引的次数
     *
     * @param[in] i
     * @return std::size_t
     */
    std::size_t degree(std::size_t i) const noexcept
    {
        if (i >= m_polys.size())
            throw std::out_of_range("ppoly: invalid index");
        return m_polys[i].degree();
    }

    /**
     * @brief 获得指定索引的多项式
     *
     * @param[in] i
     * @return const poly&
     */
    const poly<T, Real> &operator[](std::size_t i) const
    {
        if (i >= m_polys.size())
            throw std::out_of_range("ppoly: invalid index");
        return m_polys[i];
    }

    /**
     * @brief 获得指定索引的多项式引用
     *
     * @param[in] i
     * @return poly&
     */
    poly<T, Real> &operator[](std::size_t i)
    {
        if (i >= m_polys.size())
            throw std::out_of_range("ppoly: invalid index");
        return m_polys[i];
    }

    /**
     * @brief 获得传入点对应的多项式（左闭右开区间）
     *
     * @param[in] t
     * @return const poly&
     */
    const poly<T, Real> &poly_at(Real t) const
    {
        if (m_interval.empty())
            return m_polys.front();

        // 先排除超出区间范围的部分
        if (t < m_interval.front())
            return m_polys.front();
        if (t >= m_interval.back())
            return m_polys.back();
        if (t == m_interval.front())
            return m_polys[1];

        auto it = std::lower_bound(m_interval.begin(), m_interval.end(), t);
        return m_polys[it - m_interval.begin()];
    }

    /**
     * @brief 多项式求值
     *
     * @param[in] x
     * @return T
     */
    T value(Real x) const
    {
        return poly_at(x).value(x);
    }

    /**
     * @brief 在右端添加端点，增加一个零多项式
     *
     * @param[in] t
     * @return poly<T, Real>&
     */
    poly<T, Real> &push_back(Real t)
    {
        // 如果区间端点不是空，就检查是否大于最大的端点，要确保不能重合
        if (!m_interval.empty() && t < m_interval.back() + std::numeric_limits<Real>::epsilon())
            throw std::logic_error("ppoly: invalid input node");

        m_interval.emplace_back(t);
        m_polys.emplace_back(0);
        return m_polys.back();
    }

    /**
     * @brief 在左端添加端点，增加一个零多项式
     *
     * @param[in] t
     * @return poly<T, Real>&
     */
    poly<T, Real> &push_front(Real t)
    {
        if (!m_interval.empty() && t > m_interval.front() - std::numeric_limits<Real>::epsilon())
            throw std::logic_error("ppoly: invalid input node");

        m_interval.emplace(m_interval.begin(), t);
        m_polys.emplace(m_polys.begin(), 0);
        return m_polys.front();
    }

    /**
     * @brief 求导
     *
     * @return ppoly
     */
    ppoly derivative() const
    {
        auto interval = m_interval;
        ppoly ret(std::move(interval));
        for (std::size_t i = 0; i < m_polys.size(); ++i)
            ret.m_polys[i] = m_polys[i].derivative();
        return ret;
    }

    /**
     * @brief 不定积分
     *
     * @return ppoly
     */
    ppoly integrate() const
    {
        auto interval = m_interval;
        ppoly ret(std::move(interval));
        for (std::size_t i = 0; i < m_polys.size(); ++i)
            ret.m_polys[i] = m_polys[i].integrate();
        return ret;
    }

    /**
     * @brief 定积分
     *
     * @param[in] a
     * @param[in] b
     * @return T
     */
    T integrate(Real a, Real b) const
    {
        if (a >= b)
            throw std::logic_error("ppoly: invalid input interval");

        auto begin = std::lower_bound(m_interval.begin(), m_interval.end(), a) - m_interval.begin();
        auto end = std::lower_bound(m_interval.begin(), m_interval.end(), b) - m_interval.begin();

        if (begin == end)
            return m_polys[begin].integrate(a, b);

        T ret = m_polys[begin].integrate(a, m_interval[begin]);
        for (int i = begin + 1; i < end; i++)
            ret += m_polys[i].integrate(m_interval[i - 1], m_interval[i]);
        ret += m_polys[end].integrate(m_interval[end - 1], b);
        return ret;
    }

    /**
     * @brief 检查端点处的连续性
     *
     * @param[in] i
     * @return std::size_t
     */
    std::size_t continuity(std::size_t i) const
    {
        if (i >= m_interval.size())
            throw std::out_of_range("ppoly: invalid index");
        if (poly_num() == 1)
            return std::numeric_limits<std::size_t>::infinity();

        // 左右两个多项式
        poly left = m_polys[i];
        poly right = m_polys[i + 1];

        std::size_t min = std::min(left.degree(), right.degree());
        Real t = m_interval[i];

        // 标记连续性
        std::size_t c = -1;
        for (std::size_t j = 0; j < min; j++)
        {
            if (!cmp.less<Real>(left.value(t), right.value(t)))
                break;
            left = left.derivative();
            right = right.derivative();
            c++;
        }
        return c;
    }

    /**
     * @brief 整体连续性
     *
     * @return std::size_t
     */
    std::size_t continuity() const
    {
        if (poly_num() == 1)
            return std::numeric_limits<std::size_t>::infinity();

        std::size_t c = continuity(0);
        for (std::size_t i = 1; i < m_interval.size(); i++)
            c = std::min(c, continuity(i));
        return c;
    }
};

using polyf = poly<float>;
using polyd = poly<double>;
using ppolyf = ppoly<float>;
using ppolyd = ppoly<double>;

/**
 * @brief 输出函数
 *
 * @tparam T
 * @tparam Real
 * @param os
 * @param inPoly
 * @return std::ostream&
 */
template <class T, class Real> inline std::ostream &operator<<(std::ostream &os, const poly<T, Real> &inPoly)
{
    const auto &coeff = inPoly.coefficients();
    os << "f(x) = ";
    for (std::size_t i = 0; i < coeff.size(); i++)
    {
        if (i == 0)
            os << coeff[i];
        else
            os << " + " << coeff[i] << "x^" << i;
    }
    return os;
}

/**
 * @brief 流式输出
 *
 * @tparam T
 * @tparam Real
 * @param os
 * @param ppoly
 * @return std::ostream&
 */
template <class T, class Real> inline std::ostream &operator<<(std::ostream &os, const ppoly<T, Real> &ppoly)
{
    auto interval = ppoly.intervals();
    interval.emplace(interval.begin(), std::numeric_limits<Real>::lowest());
    interval.push_back(std::numeric_limits<Real>::infinity());

    os << "f(x) = {" << std::endl;
    for (std::size_t i = 0; i < ppoly.polys().size(); ++i)
    {
        os << "f_" << i << "(x) = " << ppoly.polys()[i] << " in [" << interval[i] << ", " << interval[i + 1] << "]"
           << std::endl;
    }
    os << "}";
    return os;
}

/**
 * @brief 合并两个 ppoly
 *
 * @tparam T
 * @tparam Real
 * @param p1
 * @param p2
 * @return ppoly
 */
template <class T, class Real> inline ppoly<T, Real> merge(const ppoly<T, Real> &p1, const ppoly<T, Real> &p2)
{
    if (p1.poly_num() == 1 && p2.poly_num() == 1)
        return ppoly(p1[0] + p2[0]);

    // 空多项式
    ppoly<T, Real> ret;

    // 储存所有端点，然后排序
    std::vector<Real> interval;
    interval.insert(interval.end(), p1.intervals().begin(), p1.intervals().end());
    interval.insert(interval.end(), p2.intervals().begin(), p2.intervals().end());
    std::sort(interval.begin(), interval.end());

    // 推入这些端点
    ret.push_back(interval.front());
    for (std::size_t i = 1; i < interval.size(); i++)
    {
        if (std::abs(interval[i] - interval[i - 1]) < std::numeric_limits<Real>::epsilon())
            continue;
        ret.push_back(interval[i]);
    }

    // 获取对应的多项式相加
    interval.push_back(interval.back() + 1);
    interval.emplace(interval.begin(), interval.front() - 1);
    for (std::size_t i = 0; i < ret.poly_num(); i++)
    {
        Real mid = (interval[i] + interval[i + 1]) / 2;
        ret[i] = p1.poly_at(mid) + p2.poly_at(mid);
    }
    return ret;
}

} // namespace math

} // namespace xy
